import numpy as np
import os
import struct
from array import array
# 从之前实现中导入核心函数
def conv2d(x,w,b,stride=1,padding=0):
    N,C_in,H_in,W_in = x.shape
    C_out,C_in,K_h,K_w = w.shape
    if padding > 0:
        x = np.pad(x,((0,0),(0,0),(padding,padding),(padding,padding)),'constant')
        H_in,W_in = x.shape[2],x.shape[3]
    H_out = (H_in - K_h) // stride + 1
    W_out = (W_in - K_w) // stride + 1
    output = np.zeros((N,C_out,H_out,W_out))
    for n in range(N):
        for c_out in range(C_out):
            for h in range(H_out):
                for w_idx in range(W_out):
                    h_start = h * stride
                    w_start = w_idx * stride
                    region = x[n,:,h_start:h_start+K_h,w_start:w_start+K_w]
                    output[n,c_out,h,w_idx] = np.sum(region * w[c_out]) + b[c_out]
    return output
def relu(x):
    return np.maximum(0,x)
def max_pool2d(x,kernel_size=2,stride=2):
    N,C,H,W = x.shape
    H_out = (H - kernel_size) // stride + 1
    W_out = (W - kernel_size) // stride + 1
    output = np.zeros((N,C,H_out,W_out))
    for n in range(N):
        for c in range(C):
            for h in range(H_out):
                for w in range(W_out):
                    h_start = h * stride
                    w_start = w * stride
                    region = x[n,c,h_start:h_start+kernel_size,w_start:w_start+kernel_size]
                    output[n,c,h,w] = np.max(region)
    return output
def flatten(x):
    return x.reshape(x.shape[0],-1)
def linear_layer(x,w,b):
    return np.dot(x,w) + b
# 固定随机种子，保证权重初始化一致
np.random.seed(114514)
def softmax(logits):
    # 数值稳定性处理：减去最大值避免溢出
    exp_logits = np.exp(logits - np.max(logits,axis=-1,keepdims=True))
    return exp_logits / np.sum(exp_logits,axis=-1,keepdims=True)
# MNIST数据集读取函数
def read_images(filename):
    with open(filename,'rb') as file:
        magic,size,rows,cols = struct.unpack(">IIII",file.read(16))
        if magic != 2051:
            raise ValueError('Magic number mismatch, expected 2051, got {}'.format(magic))
        image_data = array("B",file.read())
    images = []
    for i in range(size):
        img = np.array(image_data[i * rows * cols:(i + 1) * rows * cols])
        img = img.reshape(rows,cols)
        images.append(img)
    return images
class TinyCNN_for_MNIST:
    def __init__(self):
        # Conv层权重：1->4通道，3x3卷积核
        self.conv_w = np.random.randn(4,1,3,3) * 0.1
        self.conv_b = np.zeros(4)       
        # Linear层权重：14*14*4 -> 10
        self.linear_w = np.random.randn(14*14*4,10) * 0.1
        self.linear_b = np.zeros(10)  
    def forward(self,x):
        # Conv2d -> ReLU -> MaxPool2d -> Flatten -> Linear -> Softmax
        x = conv2d(x,self.conv_w,self.conv_b,stride=1,padding=1)
        x = relu(x)
        x = max_pool2d(x,kernel_size=2,stride=2)
        x = flatten(x)
        logits = linear_layer(x,self.linear_w,self.linear_b)
        probs = softmax(logits)
        return logits,probs
# 测试脚本
if __name__ == "__main__":
    # 设置MNIST测试集文件路径
    mnist_test_file = '3/mnist/t10k-images.idx3-ubyte'
    if not os.path.exists(mnist_test_file):
        print(f"错误:找不到MNIST测试集文件 '{mnist_test_file}'")
    else:
        # 加载测试图像
        test_images = read_images(mnist_test_file)
        first_test_image = test_images[0]
        # 图像预处理
        input_tensor = (first_test_image.astype(np.float32) / 255.0 - 0.5) * 2.0
        input_tensor = np.expand_dims(input_tensor,axis=(0,1))
        # 模型实例化和前向传播
        model = TinyCNN_for_MNIST()
        logits,probs = model.forward(input_tensor)
        # 结果输出
        print("Input Tensor Shape:",input_tensor.shape)
        print("Logits shape:",logits.shape,"Probs shape:",probs.shape)
        np.set_printoptions(precision=8,suppress=False)
        print("\nLogits:",logits[0])
        print("Probs:",probs[0])
        print("\nChecksum logits sum:",float(np.sum(logits)))
        print("Checksum probs sum:",float(np.sum(probs)))
